import numpy as np
import matplotlib.pyplot as plt
from pyswarm import pso
from scipy.optimize import minimize
from scipy.optimize import Bounds

# 常量定义
g = 9.8  # 重力加速度
missile_speed = 300  # 导弹速度
cloud_sink_speed = 3  # 云团下沉速度
cloud_effective_radius = 10  # 云团有效半径
cloud_duration = 20  # 云团有效时间

# 位置定义
fake_target = np.array([0, 0, 0])
real_target = np.array([0, 200, 0])  # 真目标位置
target_radius = 7  # 目标圆柱体半径
target_height = 10  # 目标圆柱体高度

M1_start = np.array([20000, 0, 2000])
FY1_start = np.array([17800, 0, 1800])

# 导弹方向向量（飞向假目标）
missile_dir = fake_target - M1_start
missile_dir_norm = np.linalg.norm(missile_dir)
missile_unit_dir = missile_dir / missile_dir_norm
missile_flight_time = missile_dir_norm / missile_speed  # 导弹总飞行时间


def missile_position(t):
    """计算导弹在t时刻的位置"""
    return M1_start + missile_unit_dir * missile_speed * t


def cloud_position(alpha, v, t_drop, t_burst, tau):
    """
    计算云团在起爆后tau时刻的位置
    alpha: 无人机飞行方向(弧度)
    v: 无人机速度
    t_drop: 投放时间(从任务开始)
    t_burst: 起爆时间(从投放开始)
    tau: 起爆后的时间
    """
    # 投放点位置
    drop_pos = FY1_start + np.array([v * np.cos(alpha) * t_drop,
                                     v * np.sin(alpha) * t_drop,
                                     0])

    # 起爆点位置（考虑重力影响）
    burst_pos = drop_pos + np.array([v * np.cos(alpha) * t_burst,
                                     v * np.sin(alpha) * t_burst,
                                     -0.5 * g * t_burst ** 2])

    # 云团位置（考虑下沉）
    return np.array([burst_pos[0],
                     burst_pos[1],
                     max(0, burst_pos[2] - cloud_sink_speed * tau)])  # 确保高度不小于0


def distance_to_line(point, line_start, line_end):
    """计算点到线段的距离"""
    line_vec = line_end - line_start
    point_vec = point - line_start
    line_len = np.linalg.norm(line_vec)

    if line_len < 1e-10:  # 避免除以零
        return np.linalg.norm(point - line_start)

    line_unitvec = line_vec / line_len
    point_vec_scaled = point_vec / line_len
    t = np.dot(line_unitvec, point_vec_scaled)

    if t < 0.0:
        t = 0.0
    elif t > 1.0:
        t = 1.0

    nearest = line_start + t * line_vec
    return np.linalg.norm(nearest - point)


def is_line_of_sight_blocked(missile_pos, cloud_pos, target_center, target_radius, target_height):
    """
    判断烟幕云团是否遮蔽导弹和目标圆柱体之间的视线
    烟幕中心到导弹与目标圆柱体边缘所有连线的距离的最大值小于10m即可有效遮蔽
    """
    # 生成目标圆柱体边缘的关键点（减少采样点以提高效率）
    # 只采样圆柱体顶部和底部的关键点
    num_points = 8  # 减少采样点数
    target_edge_points = []

    # 圆柱体底部边缘
    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_radius * np.cos(angle)
        y = target_center[1] + target_radius * np.sin(angle)
        z = target_center[2]
        target_edge_points.append(np.array([x, y, z]))

    # 圆柱体顶部边缘
    for i in range(num_points):
        angle = 2 * np.pi * i / num_points
        x = target_center[0] + target_radius * np.cos(angle)
        y = target_center[1] + target_radius * np.sin(angle)
        z = target_center[2] + target_height
        target_edge_points.append(np.array([x, y, z]))

    # 计算烟幕中心到所有导弹-目标边缘点连线的最大距离
    max_distance = 0
    for target_point in target_edge_points:
        dist = distance_to_line(cloud_pos, missile_pos, target_point)
        if dist > max_distance:
            max_distance = dist
            if max_distance > cloud_effective_radius:
                # 如果已经超过有效半径，提前返回
                return False

    # 如果最大距离小于云团有效半径，则认为视线被遮蔽
    return max_distance <= cloud_effective_radius


def effective_shielding_time(alpha, v, t_drop, t_burst):
    """
    计算有效遮蔽时间（视线被遮蔽的时间）
    """
    # 起爆时刻（从任务开始）
    t_burst_total = t_drop + t_burst

    # 云团有效时间结束时刻
    t_cloud_end = t_burst_total + cloud_duration

    # 云团有效时间不能超过导弹飞行时间
    if t_cloud_end > missile_flight_time:
        t_cloud_end = missile_flight_time

    # 增加时间步长以减少计算量
    time_step = 0.5  # 从0.1增加到0.5
    effective_time = 0

    # 检查每个时间点
    for tau in np.arange(0, min(cloud_duration, missile_flight_time - t_burst_total), time_step):
        # 当前时间（从任务开始）
        t_current = t_burst_total + tau

        # 导弹位置
        missile_pos = missile_position(t_current)

        # 云团位置
        cloud_pos = cloud_position(alpha, v, t_drop, t_burst, tau)

        # 判断视线是否被遮蔽
        if is_line_of_sight_blocked(missile_pos, cloud_pos, real_target, target_radius, target_height):
            effective_time += time_step

    return effective_time


def objective_function(x):
    """
    目标函数（最大化有效遮蔽时间）
    x = [alpha, v, t_drop, t_burst]
    """
    alpha, v, t_drop, t_burst = x

    # 约束检查
    if (v < 70 or v > 140 or
            t_drop < 0 or t_burst < 0 or
            t_drop + t_burst > missile_flight_time - cloud_duration):
        return 1e6  # 返回一个很大的值表示不可行解

    # 计算遮蔽时间（取负值因为是最小化）
    shielding_time = effective_shielding_time(alpha, v, t_drop, t_burst)

    # 如果遮蔽时间为0，添加惩罚项
    if shielding_time == 0:
        # 计算导弹和云团在起爆时刻的位置
        missile_pos_at_burst = missile_position(t_drop + t_burst)
        cloud_pos_at_burst = cloud_position(alpha, v, t_drop, t_burst, 0)

        # 计算导弹到目标的距离和云团到目标的距离
        missile_to_target_dist = np.linalg.norm(real_target - missile_pos_at_burst)
        cloud_to_target_dist = np.linalg.norm(real_target - cloud_pos_at_burst)

        # 惩罚项与距离成反比
        penalty = 1000 * (1 / (missile_to_target_dist + cloud_to_target_dist + 1))
        return penalty

    return -shielding_time


# PSO参数范围
lb = [0, 70, 0, 0]  # 下界
ub = [2 * np.pi, 140, missile_flight_time - cloud_duration, cloud_duration]  # 上界

# 设置随机种子以确保结果可重复
np.random.seed(42)

print("开始PSO优化...")
# 运行PSO优化（减少粒子数和迭代次数以提高速度）
xopt, fopt = pso(objective_function, lb, ub, swarmsize=50, maxiter=100, minstep=1e-4)

print("PSO优化完成，开始局部搜索...")
# 使用局部搜索进一步优化
bounds = Bounds(lb, ub)
result = minimize(objective_function, xopt, method='SLSQP', bounds=bounds)
xopt_refined = result.x
fopt_refined = result.fun

# 使用优化后的参数
alpha_opt, v_opt, t_drop_opt, t_burst_opt = xopt_refined


# 重新计算精确的遮蔽时间（使用更小的时间步长）
def precise_effective_shielding_time(alpha, v, t_drop, t_burst):
    """使用更小的时间步长计算精确的遮蔽时间"""
    t_burst_total = t_drop + t_burst
    t_cloud_end = t_burst_total + cloud_duration

    if t_cloud_end > missile_flight_time:
        t_cloud_end = missile_flight_time

    time_step = 0.1  # 使用更小的时间步长
    effective_time = 0

    for tau in np.arange(0, min(cloud_duration, missile_flight_time - t_burst_total), time_step):
        t_current = t_burst_total + tau
        missile_pos = missile_position(t_current)
        cloud_pos = cloud_position(alpha, v, t_drop, t_burst, tau)

        if is_line_of_sight_blocked(missile_pos, cloud_pos, real_target, target_radius, target_height):
            effective_time += time_step

    return effective_time


precise_shielding_time = precise_effective_shielding_time(alpha_opt, v_opt, t_drop_opt, t_burst_opt)

# 输出结果
print("PSO优化结果:")
print(f"飞行方向: {np.degrees(xopt[0]):.2f}°")
print(f"飞行速度: {xopt[1]:.2f} m/s")
print(f"投放时间: {xopt[2]:.2f} s")
print(f"起爆时间: {xopt[3]:.2f} s")
print(f"最大有效遮蔽时间: {-fopt:.2f} s")

print("\n局部搜索优化结果:")
print(f"飞行方向: {np.degrees(xopt_refined[0]):.2f}°")
print(f"飞行速度: {xopt_refined[1]:.2f} m/s")
print(f"投放时间: {xopt_refined[2]:.2f} s")
print(f"起爆时间: {xopt_refined[3]:.2f} s")
print(f"最大有效遮蔽时间: {-fopt_refined:.2f} s")

print(f"\n精确计算的有效遮蔽时间: {precise_shielding_time:.2f} s")

# 计算投放点和起爆点坐标
drop_point = FY1_start + np.array([v_opt * np.cos(alpha_opt) * t_drop_opt,
                                   v_opt * np.sin(alpha_opt) * t_drop_opt,
                                   0])
burst_point = drop_point + np.array([v_opt * np.cos(alpha_opt) * t_burst_opt,
                                     v_opt * np.sin(alpha_opt) * t_burst_opt,
                                     -0.5 * g * t_burst_opt ** 2])

print(f"\n投放点坐标: ({drop_point[0]:.2f}, {drop_point[1]:.2f}, {drop_point[2]:.2f})")
print(f"起爆点坐标: ({burst_point[0]:.2f}, {burst_point[1]:.2f}, {burst_point[2]:.2f})")


# 可视化遮蔽效果（可选，如果不需要可以注释掉）
def plot_shielding_effect(alpha, v, t_drop, t_burst):
    """可视化遮蔽效果"""
    t_burst_total = t_drop + t_burst
    time_points = np.linspace(t_burst_total, t_burst_total + min(cloud_duration, missile_flight_time - t_burst_total),
                              50)  # 减少点数

    遮蔽状态 = []
    for t in time_points:
        tau = t - t_burst_total
        missile_pos = missile_position(t)
        cloud_pos = cloud_position(alpha, v, t_drop, t_burst, tau)

        if is_line_of_sight_blocked(missile_pos, cloud_pos, real_target, target_radius, target_height):
            遮蔽状态.append(1)
        else:
            遮蔽状态.append(0)

    plt.figure(figsize=(10, 6))
    plt.plot(time_points, 遮蔽状态, 'b-', label='遮蔽状态')
    plt.xlabel('时间 (s)')
    plt.ylabel('遮蔽状态 (1=遮蔽, 0=未遮蔽)')
    plt.title('导弹与目标之间视线遮蔽状态随时间变化')
    plt.legend()
    plt.grid(True)
    plt.show()


print("\n生成遮蔽效果图...")
plot_shielding_effect(alpha_opt, v_opt, t_drop_opt, t_burst_opt)